import pywt
import cv2
import numpy as np

imgpath="code\\pic.png"
imgpath1="code\\pics\\USTC.png"
imgpath2="code\\pics\\DwtWater1.png"
img=cv2.imread(imgpath,0)
water=cv2.imread(imgpath1,0)
waterimg=cv2.resize(cv2.imread(imgpath2,0),(400,400))
key=3
print(type(img),type(water),img.shape,water.shape)
#arnold algorithm

def arnold(img,key):
    r=img.shape[0]
    c=img.shape[1]
    p=np.zeros((r,c),np.uint8)
    
    a,b=1,1
    for k in range(key):
        for i in range(r):
            for j in range(c):
                x=(i+b*j)%r
                y=(a*i+(a*b+1)*j)%c
                p[x,y]=img[i,j]
    return p

def inverse_arnold(img,key):
    r=img.shape[0]
    c=img.shape[1]
    p=np.zeros((r,c),np.uint8)
    a,b=1,1
    for k in range(key):
        for i in range(r):
            for j in range(c):
                x=((a*b+1)*i-b*j)%r
                y=(-a*i+j)%c
                p[x,y]=img[i,j]
    return p

#cv2.imshow("1",inverse_arnold(arnold(img,10),10))
#cv2.waitKey(0)

def setwaterMark(waterimg,img,key):
    #arnold for water image
    #cv2.imshow("original",waterimg)
    #cv2.waitKey(0)
    img=cv2.resize(img,(400,400))
    waterimg=cv2.resize(waterimg,(201,201))
    #cv2.imshow("resize",waterimg)
    #cv2.waitKey(0)
    waterimg1=arnold(waterimg,key)
    cv2.imshow("arnold for water img",waterimg1)
    cv2.waitKey(0)

    c=pywt.wavedec2(img,'db2',level=3)
    [cl,(cH3,cV3,cD3),(cH2,cV2,cD2),(cH1,cV1,cD1)]=c
    print("cl type:{0}  cl shape:  {1} ".format(type(cl),cl.shape))
    waterimg1=cv2.resize(waterimg1,(101,101))
    d=pywt.wavedec2(waterimg1,'db2',level=1)
    [cal,(ch1,cv1,cd1)]=d
    print("cal type:{0}  cal shape:  {1} ".format(type(cal),cal.shape))
    a1,a2,a3,a4=0.1,0.1,0.1,0.1
    cl=cl+cal*a1
    cH3=cH3+ch1*a2
    cV3=cV3+cv1*a3
    cD3=cD3+cd1*a4

    #remake img
    newImg=pywt.waverec2([cl,(cH3,cV3,cD3),(cH2,cV2,cD2),(cH1,cV1,cD1)],'db2')
    newImg=np.array(newImg,np.uint8)
    #done 
    cv2.imwrite("code\\pics\\DwtWater.png",newImg)
    cv2.imshow("recover",newImg)
    cv2.waitKey(0)

def getWaterMark(original,waterimg,key):
    original=cv2.resize(original,(400,400))
    c=pywt.wavedec2(waterimg,'db2',level=3)
    [cl,(cH3,cV3,cD3),(cH2,cV2,cD2),(cH1,cV1,cD1)]=c
    d=pywt.wavedec2(original,'db2',level=3)
    print("cl.shape: {0}  dl.type:  {1}".format(cl.shape,len(d[0])))
    [dl,(dH3,dV3,dD3),(dH2,dV2,dD2),(dH1,dV1,dD1)]=d
    a1,a2,a3,a4=0.1,0.1,0.1,0.1
    cal=(cl-dl)/a1
    ch1=(cH3-dH3)/a2
    cv1=(cV3-dV3)/a3
    cd1=(cD3-dD3)/a4

    waterImg=pywt.waverec2([cal,(ch1,cv1,cd1)],'db2')
    waterImg=np.array(waterImg,np.uint8)
    cv2.imshow("get",waterImg)
    cv2.waitKey(0)

    waterImg=inverse_arnold(waterImg,key)
    print("extract watermarking successfully!")
    cv2.imshow("get",waterImg)
    cv2.imwrite("code\\pics\\DwtGetResult1.png",waterImg)
    cv2.waitKey(0)

setwaterMark(water,img,key)
getWaterMark(img,waterimg,key)